# -*-coding:utf-8 -*-

import time
import os
import json
import traceback
from urllib.parse import quote_plus
from obs import ObsClient
from obs import ACL
from obs import HeadPermission
from huaweicloudsdkcore.auth.credentials import BasicCredentials
from huaweicloudsdkmoderation.v3.region.moderation_region import ModerationRegion
from huaweicloudsdkcore.exceptions import exceptions
from huaweicloudsdkmoderation.v3 import *
from video_moderation_unzip import FileProcessing
from urllib.parse import unquote_plus

currentDir = os.environ["RUNTIME_CODE_ROOT"]
default_video = currentDir + '/data/default.png'


def check_configuration(context):
    region = context.getUserData('region')
    if not region:
        return 'region is not configured'
    obs_server = context.getUserData('obs_server')
    if not obs_server:
        return 'obs_server is not configured'
    ak = context.getAccessKey().strip()
    sk = context.getSecretKey().strip()
    if not ak or not sk:
        ak = context.getUserData('ak', '').strip()
        sk = context.getUserData('sk', '').strip()
        if not ak or not sk:
            return 'ak or sk is empty'


def handler(event, context):
    log = context.getLogger()
    result = check_configuration(context)
    if result is not None:
        return result
    records = event.get("Records", None)
    if records is None:
        return 'Records is empty'

    processor = Processor(context)
    file_processing = FileProcessing(context)

    try:
        for record in records:
            bucket_name, object_key = get_obs_obj_info(record)
            if object_key.endswith("zip") or object_key.endswith("tar"):
                file_processing.run(record)
                records = event.get("Records", None)
                for record in records:
                    bucket_name, object_key = get_obs_obj_info(record)
                    if object_key.endswith("zip") or object_key.endswith("tar"):
                        break
                    processor.process(record)
            elif object_key.endswith("png"):
                break
            else:
                processor.process(record)
    except exceptions as e:
        log.error("failed to process video, "
                  f"exception: {traceback.format_exc()}")
    finally:
        processor.obs_client.close()
    return 'Done'


class Processor:
    def __init__(self, context=None):
        self.log = context.getLogger()
        self.obs_client = new_obs_client(context)
        self.moderation_client = new_moderation_client(context)
        self.dump_bucket = context.getUserData('dump_bucket')
        self.video_url = None
        self.src_bucket = None
        self.src_object_key = None
        self.log_object_key = None

    def process(self, record):
        self.parse_record(record)
        self.set_object_acl()
        response_create = self.create_video_moderation()
        if response_create is None:
            self.log.error("failed to run video moderation")
            return

        while True:
            response = self.query_video_moderation(response_create.job_id)
            if response.status != "running":
                break
            else:
                time.sleep(10)

        result = response.result
        if result.suggestion == "pass":
            return
        elif result.suggestion == "block":
            self.dump_video()
            self.delete_video()
            self.upload_replacement_video()
            self.report_to_obs(result)
        elif result.suggestion == "review":
            self.report_to_obs(result)
        else:
            self.log.warn("suggestion %s is not in (pass, block, review)",
                          result.suggestion)

    def parse_record(self, record):
        region = get_region(record)
        (bucket_name, object_key) = get_obs_obj_info(record)
        self.log.info("input bucket_name: %s", bucket_name)
        object_key = unquote_plus(object_key)
        self.log.info("input object: %s", object_key)
        self.src_bucket = bucket_name
        self.src_object_key = object_key

        self.video_url = "https://" + bucket_name + ".obs." + \
                         region + ".myhuaweicloud.com/" + quote_plus(object_key)
        (path, filename) = os.path.split(object_key)
        (filename, _) = os.path.splitext(filename)

        self.log_object_key = path + unquote_plus(filename) + ".json"

    def set_object_acl(self):
        try:
            resp = self.obs_client.setObjectAcl(self.src_bucket, self.src_object_key,
                                                aclControl=HeadPermission.PUBLIC_READ)
            if resp.status < 300:
                self.log.info("succeed to set object acl.")
            else:
                self.log.error("failed to set object acl, "
                               f"requestId: {resp.requestId} "
                               f"errorCode: {resp.errorCode} "
                               f"errorMessage: {resp.errorMessage}")
        except exceptions as e:
            self.log.error("failed to set object acl, ", traceback.format_exc())

    def create_video_moderation(self):
        try:
            request = RunCreateVideoModerationJobRequest()
            listAudioCategoriesbody = [
                "porn",
                "ad",
                "politics",
                "moan",
                "abuse"
            ]
            listImageCategoriesbody = [
                "politics",
                "image_text",
                "porn",
                "terrorism"
            ]
            databody = VideoCreateRequestData(
                url=self.video_url
            )
            request.body = VideoCreateRequest(
                audio_categories=listAudioCategoriesbody,
                image_categories=listImageCategoriesbody,
                event_type="default",
                data=databody
            )
            response = self.moderation_client.run_create_video_moderation_job(request)
            return response
        except exceptions.ClientRequestException as e:
            self.log.error(f"failed to run create video moderation, "
                           f"status_code:{e.status_code}, "
                           f"request_id:{e.request_id}, "
                           f"error_code:{e.error_code}. "
                           f"error_msg:{e.error_msg}")

    def query_video_moderation(self, job_id):
        try:
            request = RunQueryVideoModerationJobRequest()
            request.job_id = job_id
            response = self.moderation_client.run_query_video_moderation_job(request)
            return response
        except exceptions.ClientRequestException as e:
            self.log.error(f"failed to run query video moderation, "
                           f"status_code:{e.status_code}, "
                           f"request_id:{e.request_id}, "
                           f"error_code:{e.error_code}. "
                           f"error_msg:{e.error_msg}")

    def dump_video(self):
        try:
            resp = self.obs_client.copyObject(self.src_bucket,
                                              self.src_object_key,
                                              self.dump_bucket,
                                              self.src_object_key)
            if resp.status < 300:
                self.log.info("succeed to dump video.")
            else:
                self.log.error("failed to dump video, "
                               f"requestId: {resp.requestId} "
                               f"errorCode: {resp.errorCode} "
                               f"errorMessage: {resp.errorMessage}")
        except:
            self.log.error("failed to dump video, "
                           f"exception: {traceback.format_exc()}")

    def delete_video(self):
        try:
            resp = self.obs_client.deleteObject(self.src_bucket, self.src_object_key)

            if resp.status < 300:
                self.log.info("succeed to delete video.")
            else:
                self.log.error("failed to delete video, "
                               f"requestId: {resp.requestId} "
                               f"errorCode: {resp.errorCode} "
                               f"errorMessage: {resp.errorMessage}")
        except exceptions as e:
            self.log.error("failed to delete video, ", traceback.format_exc())

    def upload_replacement_video(self):
        try:
            name, ext = os.path.splitext(self.src_object_key)
            self.src_object_key = name + ".png"
            resp = self.obs_client.putFile(self.src_bucket, self.src_object_key,
                                           default_video)
            if resp.status < 300:
                self.log.info("succeed to upload replacement video.")
            else:
                self.log.error("failed to upload replacement video, "
                               f"requestId: {resp.requestId} "
                               f"errorCode: {resp.errorCode} "
                               f"errorMessage: {resp.errorMessage}")
        except exceptions as e:
            self.log.error("failed to replace video, "
                           f"exception: {traceback.format_exc()}")

    def report_to_obs(self, result):
        labels = {
            "moderation_result": result.suggestion,
            "original_path": "obs://" + self.src_bucket + "/" +
                             self.src_object_key,
            "dump_path": "obs://" + self.dump_bucket + "/" +
                         self.src_object_key,
        }
        payload = {"log_time_ns": int(round(time.time() * 1000000000)),
                   "contents": str(result), "labels": labels}
        resp = self.obs_client.putObject(self.dump_bucket, self.log_object_key,
                                         json.dumps(payload))
        if resp.status < 300:
            self.log.info("succeed to upload log to obs.")
        else:
            self.log.error('failed to upload log to obs, errorCode: %s, '
                           'errorMessage: %s' % (resp.errorCode,
                                                 resp.errorMessage))
            return False


def get_obs_obj_info(record):
    s3 = record['s3']
    return (s3['bucket']['name'], s3['object']['key'])


def new_obs_client(context):
    return ObsClient(
        access_key_id=context.getAccessKey(),
        secret_access_key=context.getSecretKey(),
        server=context.getUserData('obs_server'),
    )


def new_moderation_client(context):
    credentials = BasicCredentials(context.getAccessKey(),
                                   context.getSecretKey())
    return ModerationClient.new_builder() \
        .with_credentials(credentials) \
        .with_region(ModerationRegion.value_of(context.getUserData('region'))) \
        .build()


def get_region(record):
    if 'eventRegion' in record:
        return record.get("eventRegion")
    else:
        return record.get("awsRegion")